package com.example.state;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

/**
 * @author lbq
 * @version 1.0
 * @date 2022/3/27 10:22
 */
public class SumStateDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        final DataStreamSource<Tuple2<Integer, Integer>> data = env.fromElements(
                Tuple2.of(1, 3),
                Tuple2.of(1, 5),
                Tuple2.of(1, 7),
                Tuple2.of(1, 4),
                Tuple2.of(1, 2),
                Tuple2.of(2, 3),
                Tuple2.of(2, 5),
                Tuple2.of(2, 7),
                Tuple2.of(2, 4),
                Tuple2.of(2, 2)
        );

        final KeyedStream<Tuple2<Integer, Integer>, Integer> keyedStream = data.keyBy(value -> value.f0);

        keyedStream.flatMap(new RichFlatMapFunction<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>>() {

            private ValueState<Tuple2<Integer, Integer>> sum;

            @Override
            public void flatMap(Tuple2<Integer, Integer> value, Collector<Tuple2<Integer, Integer>> out) throws Exception {
                final Tuple2<Integer, Integer> currentState = sum.value();
                currentState.f0 += 1;
                currentState.f1 += value.f1;
                System.out.println("currentState = " + currentState);
                sum.update(currentState);

                if (currentState.f0 >= 2) {
                    out.collect(new Tuple2<>(value.f0, currentState.f1 / currentState.f0));
                    // sum.clear();
                }
            }

            @Override
            public void open(Configuration parameters) throws Exception {
                final ValueStateDescriptor<Tuple2<Integer, Integer>> descriptor = new ValueStateDescriptor<>("average", TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {
                }), Tuple2.of(0, 0));
                sum = getRuntimeContext().getState(descriptor);
            }
        }).print();

        env.execute();
    }
}
